from Crypto.PublicKey import RSA

#import Crypto.Cipher.PKCS1_v1_5 as cipher
import Crypto.Signature.PKCS1_v1_5 as signer

import Crypto.Cipher.PKCS1_OAEP as cipher
#import Crypto.Signature.PKCS1_PSS as signer

from Crypto import Hash
import time
import pandas as pd


'''
cipher_obj = cipher.new(RSA.importKey(pk))
cipher_text = cipher_obj.encrypt(message.encode())

cipher_obj = cipher.new(RSA.importKey(sk))
plaint_text = cipher_obj.decrypt(cipher_text, 'DecryptError').decode() # for pkcs
#cipher_obj.decrypt(cipher_text.encode()).decode() # for oaep

print(message == plaint_text)

# 写入文件--1
with open("c.pem", "wb") as x:
   x.write(s_key)
with open("d.pem", "wb") as x:
   x.write(g_key)

with open('c.pem','rb')as x:
   s_key = RSA.importKey(x.read())
 # new_g_key = s_key.publickey().export_key()
 # print(new_g_key)
'''


def generate_key(bits):
    return RSA.generate(bits)


def encrypt(message, pk, hasher=Hash.SHA1):
    cipher_obj = cipher.new(RSA.importKey(pk))
    org_bytes = message.encode()
    #length_en = RSA.RsaKey.size_in_bytes(RSA.importKey(pk)) - 11 # 2048/8-11=245
    length_en = RSA.RsaKey.size_in_bytes(RSA.importKey(pk)) - 2 * hasher.digest_size - 2  # 2048/8-(2*20+2) -- SHA1

    res_en = b''
    for i in range(0, len(org_bytes), length_en):
        res_en += cipher_obj.encrypt(org_bytes[i: i + length_en])
    cipher_text = res_en
    return cipher_text


def decrypt(cipher_text, sk):
    cipher_obj = cipher.new(RSA.importKey(sk))
    length_de = RSA.RsaKey.size_in_bytes(RSA.importKey(sk))

    res_de = b''
    for i in range(0, len(cipher_text), length_de):
        #res_de += cipher_obj.decrypt(cipher_text[i:i + length_de], 'DecryptError') # for pkcs
        res_de += cipher_obj.decrypt(cipher_text[i:i + length_de])  # for oaep
    plaint_text = res_de.decode()
    return plaint_text



def sign(plaint_text, sk):
    signer_obj = signer.new(RSA.importKey(sk))
    hash_obj = Hash.SHA256.new()
    hash_obj.update(plaint_text.encode())
    signature = signer_obj.sign(hash_obj)
    return signature

def verify(signature, plaint_text, pk):
    verify_obj = signer.new(RSA.importKey(pk))
    hash_obj = Hash.SHA256.new()
    hash_obj.update(plaint_text.encode())
    verify = verify_obj.verify(hash_obj, signature)
    return verify


if __name__ == "__main__":
    key = generate_key(2048)
    pk = key.publickey().export_key()
    sk = key.export_key()

    data = []
    for i in range(100, 10100, 100):
        start = time.time()

        message = "We are different, work hard!" * i
        """
        cipher_text = encrypt(message, pk)
        plaint_text = decrypt(cipher_text, sk)
        print(message == plaint_text)
        """

        signature = sign(message, sk)
        flag = verify(signature, message, pk)
        print(flag)

        end = time.time()
        total = end - start

        data.append([i, total])
    pd.DataFrame(data, columns=["count", "time"]).to_csv("sign_pkcs.csv", index=False, encoding='utf-8')
